package linker

import (
	"debug/elf"
	"math"
	"math/bits"
	"risc-v/ld/pkg/utils"
)

// InputSection 表示的结构和 SectionHeader 所指向的数据相对应
// 可以理解为 InputSection 用来表示某一个 Section 实体中的信息
type InputSection struct {
	File               *ObjectFile //所属的 ObjectFile 文件 即.o 文件
	Contents           []byte      //此 InputSection 的内容，也就是 ELF文件中的数据块
	SectionHeaderIndex uint32      //此 InputSection 所对应的 SectionHeader 在ELF文件中的 SectionHeaderTable 中的索引
	SectionSize        uint32      // InputSection 所属的 Section 所对应的size
	IsAlive            bool        //是否要保存到最终de可执行文件中
	P2Align            uint8       // 2 的幂数

	Offset        uint32
	OutPutSection *OutPutSection // 该 InputSection 所属的 OutPutSection

	RelSectionIndex uint32 // 重定位节的索引
	Rels            []Rela // 所有的重定位节
}

// NewInputSection 从 ObjectFile 实际对应的数据块来创建一个  InputSection 实例
// file :所属的 ObjectFile
// sectionHeaderIndex: InputSection 对应的 SectionHeader 在 SectionHeaderTable 中的索引(下标)
func NewInputSection(ctx *Context, name string, file *ObjectFile, sectionHeaderIndex uint32) *InputSection {
	s := &InputSection{
		File:               file,
		SectionHeaderIndex: sectionHeaderIndex,
		IsAlive:            true, // 默认是 true 最终用不到不会设置为 false
		Offset:             math.MaxUint32,
		RelSectionIndex:    math.MaxUint32,
		SectionSize:        math.MaxUint32,
	}
	sectionHeader := s.SectionHeader()
	// 读取 Section 内容
	s.Contents = file.File.Content[sectionHeader.Offset : sectionHeader.Offset+sectionHeader.Size]

	// 断言 section 的类型 这里不处理 被压缩的节
	utils.Assert(sectionHeader.Flags&uint64(elf.SHF_COMPRESSED) == 0)
	s.SectionSize = uint32(sectionHeader.Size)

	toP2Align := func(align uint64) uint8 {
		if align == 0 {
			return 0
		}
		// align 的二进制后面有几个零
		return uint8(bits.TrailingZeros64(align))
	}
	// 初始化对齐变量
	s.P2Align = toP2Align(sectionHeader.AddrAlign)

	// OutPutSection 的初始化
	s.OutPutSection = GetOutPutSection(ctx, name, uint64(sectionHeader.Type), sectionHeader.Flags)

	return s
}

// SectionHeader 返回所属的 InputSection 对应的 SectionHeader
func (i *InputSection) SectionHeader() *SectionHeader {
	// 确保 SectionHeaderIndex 未越界
	utils.Assert(i.SectionHeaderIndex < uint32(len(i.File.Sections)))
	return &i.File.SectionHeaders[i.SectionHeaderIndex]
}

// Name 返回此 InputSection 所对应的 name
// 内部主要通过 GetSectionEntryName() 方法实现
// 从字符串表中获取真实
func (i *InputSection) Name() string {
	return GetSectionEntryName(i.File.SectionHeaderStrTable, i.SectionHeader().Name)
}

func (i *InputSection) WriteTo(ctx *Context, buf []byte) {
	// 如何节的类型不占用内存 不分配 如 .bss 段
	// 不占用磁盘空间 但是在运行时占用内存空间
	if i.SectionHeader().Type == uint32(elf.SHT_NOBITS) || i.SectionSize == 0 {
		return
	}
	i.CopyContents(buf)

	if i.SectionHeader().Flags&uint64(elf.SHF_ALLOC) != 0 {
		//TODO
		i.ApplyRelocateAlloc(ctx, buf)
	}
}

// CopyContents 拷贝 InputSection 中的 Contents 到 bug
func (i *InputSection) CopyContents(buf []byte) {
	copy(buf, i.Contents)
}

// GetRelSection 获取重定位节
// 从数据中读入重定位节
func (i *InputSection) GetRelSection() []Rela {
	if i.RelSectionIndex == math.MaxUint32 || i.Rels != nil {
		return i.Rels
	}

	byteStream := i.File.GetBytesFromSectionHeader(&i.File.InputFile.SectionHeaders[i.RelSectionIndex])
	i.Rels = utils.ReadSlice[Rela](byteStream, RelaSize)

	return i.Rels
}

// GetAddr 返回 Section 在输出文件中的地址
// 即对应的 OutPutSection 的地址 + InputSection.Offset (节偏移)
func (i *InputSection) GetAddr() uint64 {
	return i.OutPutSection.SectionHeader.Addr + uint64(i.Offset)
}

func (i *InputSection) ScanRelocations() {
	for _, rel := range i.GetRelSection() {
		symbol := i.File.Symbols[rel.Symbol]
		if symbol.File == nil {
			continue
		}

		if rel.Type == uint32(elf.R_RISCV_TLS_GOT_HI20) {
			symbol.Flags |= NeedGotYp
		}
	}
}

func (i *InputSection) ApplyRelocateAlloc(ctx *Context, base []byte) {
	rels := i.GetRelSection()
	for a := 0; a < len(rels); a++ {
		rel := rels[a]
		if rel.Type == uint32(elf.R_RISCV_NONE) ||
			rel.Type == uint32(elf.R_RISCV_RELAX) {
			continue
		}

		sym := i.File.Symbols[rel.Symbol]
		loc := base[rel.Offset:]

		if sym.File == nil {
			continue
		}

		S := sym.GetAddr() //的到符号的真实地址
		A := uint64(rel.Addend)
		P := i.GetAddr() + rel.Offset

		switch elf.R_RISCV(rel.Type) {
		case elf.R_RISCV_32:
			utils.Write[uint32](loc, uint32(S+A))
		case elf.R_RISCV_64:
			utils.Write[uint64](loc, S+A)
		case elf.R_RISCV_BRANCH:
			writeBtype(loc, uint32(S+A-P))
		case elf.R_RISCV_JAL:
			writeJtype(loc, uint32(S+A-P))
		case elf.R_RISCV_CALL, elf.R_RISCV_CALL_PLT:
			val := uint32(S + A - P)
			writeUtype(loc, val)
			writeItype(loc[4:], val)
		case elf.R_RISCV_TLS_GOT_HI20:
			utils.Write[uint32](loc, uint32(sym.GetGotTpAddr(ctx)+A-P))
		case elf.R_RISCV_PCREL_HI20:
			utils.Write[uint32](loc, uint32(S+A-P))
		case elf.R_RISCV_HI20:
			writeUtype(loc, uint32(S+A))
		case elf.R_RISCV_LO12_I, elf.R_RISCV_LO12_S:
			val := S + A
			if rel.Type == uint32(elf.R_RISCV_LO12_I) {
				writeItype(loc, uint32(val))
			} else {
				writeStype(loc, uint32(val))
			}

			if utils.SignExtend(val, 11) == val {
				setRs1(loc, 0)
			}
		case elf.R_RISCV_TPREL_LO12_I, elf.R_RISCV_TPREL_LO12_S:
			val := S + A - ctx.TpAddr
			if rel.Type == uint32(elf.R_RISCV_TPREL_LO12_I) {
				writeItype(loc, uint32(val))
			} else {
				writeStype(loc, uint32(val))
			}

			if utils.SignExtend(val, 11) == val {
				setRs1(loc, 4)
			}
		}
	}

	//
	for a := 0; a < len(rels); a++ {
		switch elf.R_RISCV(rels[a].Type) {
		case elf.R_RISCV_PCREL_LO12_I, elf.R_RISCV_PCREL_LO12_S:
			sym := i.File.Symbols[rels[a].Symbol]
			utils.Assert(sym.InputSection == i)
			loc := base[rels[a].Offset:]
			val := utils.Read[uint32](base[sym.Value:])

			if rels[a].Type == uint32(elf.R_RISCV_PCREL_LO12_I) {
				writeItype(loc, val)
			} else {
				writeStype(loc, val)
			}
		}
	}

	for a := 0; a < len(rels); a++ {
		switch elf.R_RISCV(rels[a].Type) {
		case elf.R_RISCV_PCREL_HI20, elf.R_RISCV_TLS_GOT_HI20:
			loc := base[rels[a].Offset:]
			val := utils.Read[uint32](loc)
			utils.Write[uint32](loc, utils.Read[uint32](i.Contents[rels[a].Offset:]))
			writeUtype(loc, val)
		}
	}
}

func itype(val uint32) uint32 {
	return val << 20
}

func stype(val uint32) uint32 {
	return utils.Bits(val, 11, 5)<<25 | utils.Bits(val, 4, 0)<<7
}

func btype(val uint32) uint32 {
	return utils.Bit(val, 12)<<31 | utils.Bits(val, 10, 5)<<25 |
		utils.Bits(val, 4, 1)<<8 | utils.Bit(val, 11)<<7
}

func utype(val uint32) uint32 {
	return (val + 0x800) & 0xffff_f000
}

func jtype(val uint32) uint32 {
	return utils.Bit(val, 20)<<31 | utils.Bits(val, 10, 1)<<21 |
		utils.Bit(val, 11)<<20 | utils.Bits(val, 19, 12)<<12
}

func cbtype(val uint16) uint16 {
	return utils.Bit(val, 8)<<12 | utils.Bit(val, 4)<<11 | utils.Bit(val, 3)<<10 |
		utils.Bit(val, 7)<<6 | utils.Bit(val, 6)<<5 | utils.Bit(val, 2)<<4 |
		utils.Bit(val, 1)<<3 | utils.Bit(val, 5)<<2
}

func cjtype(val uint16) uint16 {
	return utils.Bit(val, 11)<<12 | utils.Bit(val, 4)<<11 | utils.Bit(val, 9)<<10 |
		utils.Bit(val, 8)<<9 | utils.Bit(val, 10)<<8 | utils.Bit(val, 6)<<7 |
		utils.Bit(val, 7)<<6 | utils.Bit(val, 3)<<5 | utils.Bit(val, 2)<<4 |
		utils.Bit(val, 1)<<3 | utils.Bit(val, 5)<<2
}

func writeItype(loc []byte, val uint32) {
	mask := uint32(0b000000_00000_11111_111_11111_1111111)
	utils.Write[uint32](loc, (utils.Read[uint32](loc)&mask)|itype(val))
}

func writeStype(loc []byte, val uint32) {
	mask := uint32(0b000000_11111_11111_111_00000_1111111)
	utils.Write[uint32](loc, (utils.Read[uint32](loc)&mask)|stype(val))
}

func writeBtype(loc []byte, val uint32) {
	mask := uint32(0b000000_11111_11111_111_00000_1111111)
	utils.Write[uint32](loc, (utils.Read[uint32](loc)&mask)|btype(val))
}

func writeUtype(loc []byte, val uint32) {
	mask := uint32(0b000000_00000_00000_000_11111_1111111)
	utils.Write[uint32](loc, (utils.Read[uint32](loc)&mask)|utype(val))
}

func writeJtype(loc []byte, val uint32) {
	mask := uint32(0b000000_00000_00000_000_11111_1111111)
	utils.Write[uint32](loc, (utils.Read[uint32](loc)&mask)|jtype(val))
}

func setRs1(loc []byte, rs1 uint32) {
	utils.Write[uint32](loc, utils.Read[uint32](loc)&0b111111_11111_00000_111_11111_1111111)
	utils.Write[uint32](loc, utils.Read[uint32](loc)|(rs1<<15))
}
